import requests
import json
from typing import Dict, List, Optional
from datetime import datetime, timedelta
import time

class TrainAPI:
    def __init__(self):
        self.base_url = "https://kyfw.12306.cn/otn"
        self.headers = {
            'User-Agent': 'Mozilla/5.0 (Windows NT 10.0; Win64; x64) AppleWebKit/537.36',
            'Accept': 'application/json, text/javascript, */*; q=0.01',
            'Accept-Language': 'zh-CN,zh;q=0.9,en;q=0.8',
            'Accept-Encoding': 'gzip, deflate, br',
            'Connection': 'keep-alive',
        }
        
    def get_stations(self) -> Dict[str, str]:
        """获取所有火车站信息"""
        url = f"{self.base_url}/resources/js/framework/station_name.js"
        
        try:
            response = requests.get(url, headers=self.headers, timeout=10)
            if response.status_code == 200:
                # 解析12306的站点数据格式
                content = response.text
                # 提取站点信息（简化版本）
                stations = {}
                # 这里需要解析12306的特殊格式
                return self._parse_stations(content)
        except Exception as e:
            print(f"获取火车站信息失败: {e}")
            
        return {}
    
    def _parse_stations(self, content: str) -> Dict[str, str]:
        """解析12306站点数据"""
        stations = {}
        try:
            # 12306的站点数据格式：@bjb|北京北|VAP|beijingbei|bjb|0
            lines = content.split('@')
            for line in lines[1:]:  # 跳过第一个空元素
                parts = line.split('|')
                if len(parts) >= 6:
                    code = parts[0]
                    name = parts[1]
                    stations[name] = code
        except Exception as e:
            print(f"解析站点数据失败: {e}")
            
        return stations
    
    def search_trains(self, from_station: str, to_station: str, 
                     date: str = None) -> List[Dict]:
        """搜索火车票"""
        if not date:
            date = (datetime.now() + timedelta(days=1)).strftime('%Y-%m-%d')
            
        # 这里使用模拟数据，实际需要调用12306的查询接口
        # 由于12306有反爬虫机制，这里提供模拟数据
        return self._get_mock_train_data(from_station, to_station, date)
    
    def _get_mock_train_data(self, from_station: str, to_station: str, date: str) -> List[Dict]:
        """模拟火车票数据"""
        mock_trains = [
            {
                'train_no': 'G1',
                'train_type': '高铁',
                'from_station': from_station,
                'to_station': to_station,
                'departure_time': '06:00',
                'arrival_time': '11:30',
                'duration': '5小时30分',
                'price': {
                    'business': 1748,
                    'first': 924,
                    'second': 553
                },
                'seats': {
                    'business': '有',
                    'first': '有',
                    'second': '有'
                },
                'date': date
            },
            {
                'train_no': 'G3',
                'train_type': '高铁',
                'from_station': from_station,
                'to_station': to_station,
                'departure_time': '08:00',
                'arrival_time': '13:30',
                'duration': '5小时30分',
                'price': {
                    'business': 1748,
                    'first': 924,
                    'second': 553
                },
                'seats': {
                    'business': '有',
                    'first': '有',
                    'second': '有'
                },
                'date': date
            },
            {
                'train_no': 'D1',
                'train_type': '动车',
                'from_station': from_station,
                'to_station': to_station,
                'departure_time': '07:30',
                'arrival_time': '14:00',
                'duration': '6小时30分',
                'price': {
                    'first': 924,
                    'second': 553
                },
                'seats': {
                    'first': '有',
                    'second': '有'
                },
                'date': date
            }
        ]
        
        return mock_trains
    
    def get_train_schedule(self, train_no: str, date: str) -> List[Dict]:
        """获取列车时刻表"""
        # 模拟时刻表数据
        mock_schedule = [
            {
                'station_name': '北京南',
                'arrival_time': '--',
                'departure_time': '06:00',
                'stop_time': '--',
                'distance': 0
            },
            {
                'station_name': '天津南',
                'arrival_time': '06:35',
                'departure_time': '06:37',
                'stop_time': '2分',
                'distance': 120
            },
            {
                'station_name': '济南西',
                'arrival_time': '08:30',
                'departure_time': '08:32',
                'stop_time': '2分',
                'distance': 400
            },
            {
                'station_name': '南京南',
                'arrival_time': '10:45',
                'departure_time': '10:47',
                'stop_time': '2分',
                'distance': 800
            },
            {
                'station_name': '上海虹桥',
                'arrival_time': '11:30',
                'departure_time': '--',
                'stop_time': '--',
                'distance': 1200
            }
        ]
        
        return mock_schedule
    
    def get_ticket_price(self, train_no: str, from_station: str, 
                        to_station: str, seat_type: str) -> Dict:
        """获取票价信息"""
        # 模拟票价数据
        base_prices = {
            'G': {'business': 1748, 'first': 924, 'second': 553},
            'D': {'first': 924, 'second': 553},
            'K': {'hard_sleeper': 300, 'soft_sleeper': 450, 'hard_seat': 150}
        }
        
        train_type = train_no[0]
        if train_type in base_prices:
            return base_prices[train_type].get(seat_type, 0)
        
        return 0